package xyz.sundy.sftp.client;

import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;

import java.util.Map;

/**
 * @author Sundy
 * @since 2025/7/2 17:07
 */
@Slf4j
public class SftpPoolObjectFactory extends BasePooledObjectFactory<ChannelSftp> {
    private final String sftpName;
    private final Map<String, String> sftpConfig;

    public SftpPoolObjectFactory(String sftpName, Map<String, String> sftpConfig) {
        this.sftpName = sftpName;
        this.sftpConfig = sftpConfig;
    }

    @Override
    public ChannelSftp create() throws Exception {
        String host = sftpConfig.get("sftp_host");
        int port = Integer.parseInt(sftpConfig.get("sftp_port"));
        String userName = sftpConfig.get("sftp_user");

        if (sftpConfig.containsKey("sftp_pwd")) {
            String password = sftpConfig.get("sftp_pwd");
            return getConnectionByPwd(host, port, userName, password);
        } else if (sftpConfig.containsKey("sftp_key")) {
            String keyFile = sftpConfig.get("sftp_key");
            return getConnectionByKey(host, port, userName, keyFile);
        } else {
            log.info("SFTP[{}] 配置文件中没有配置密码和秘钥", sftpName);
            throw new SftpConnectionError("SFTP[" + sftpName + "] 配置文件中没有配置密码和秘钥");
        }
    }

    @Override
    public PooledObject<ChannelSftp> wrap(ChannelSftp obj) {
        return new DefaultPooledObject<>(obj);
    }

    @Override
    public void destroyObject(PooledObject<ChannelSftp> p) throws Exception {
        ChannelSftp sftpHandler = p.getObject();
        if (sftpHandler.isConnected()) {
            sftpHandler.disconnect();
        }
        Session session = sftpHandler.getSession();
        if (session != null && session.isConnected()) {
            session.disconnect();
        }
    }

    @Override
    public boolean validateObject(PooledObject<ChannelSftp> p) {
        ChannelSftp sftpHandler = p.getObject();
        if( sftpHandler.isConnected()) {
            try {
                sftpHandler.pwd();
                return true;
            } catch (SftpException e) {
                // 连接异常，重新连接
            }
        };
        return false;
    }

    // 通过密码创建 SFTP 连接
    private ChannelSftp getConnectionByPwd(String host, int port, String userName, String password) throws JSchException {
        JSch jsch = new JSch();
        Session session = jsch.getSession(userName, host, port);
        session.setPassword(password);
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect(60000);
        Channel channel = session.openChannel("sftp");
        channel.connect();
        return (ChannelSftp) channel;
    }

    // 通过密钥创建 SFTP 连接
    private ChannelSftp getConnectionByKey(String host, int port, String userName, String keyFile) throws JSchException {
        JSch jsch = new JSch();
        jsch.addIdentity(keyFile);
        Session session = jsch.getSession(userName, host, port);
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect(60000);
        Channel channel = session.openChannel("sftp");
        channel.connect();
        return (ChannelSftp) channel;
    }
}